package org.basis.algorithm.promote;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Stack;

/**
 * 并查集
 * 用于快速的合并两个集合，以及快速的查询两个元素是否属于同一个集合，使得时间复杂度达到O（1）
 *
 * @author Mr_wenpan@163.com 2022/01/05 15:22
 */
public class UnionFind {

    /**
     * 元素，将用户传入的v包裹一层
     */
    static class Element<V> {

        V value;

        public Element(V value) {
            this.value = value;
        }
    }

    /**
     * 并查集
     */
    public static class UnionFindSet<V> {
        /**
         * 存放用户节点和包裹节点的对应关系
         */
        Map<V, Element<V>> elementMap = new HashMap<>();
        /**
         * 存放包裹节点的父节点
         */
        Map<Element<V>, Element<V>> fatherMap = new HashMap<>();
        /**
         * 存放头节点对应下面的元素个数
         */
        Map<Element<V>, Integer> sizeMap = new HashMap<>();

        /**
         * 并查集要求用户在使用的时候就将元素给定
         */
        public UnionFindSet(List<V> list) {
            for (V v : list) {
                Element<V> element = new Element<>(v);
                elementMap.put(v, element);
                fatherMap.put(element, element);
                sizeMap.put(element, 1);
            }
        }

        /**
         * 合并ab两个元素所在的集合
         */
        public void union(V a, V b) {
            if (!elementMap.containsKey(a) || !elementMap.containsKey(b)) {
                return;
            }
            Element<V> head1 = findHead(a);
            Element<V> head2 = findHead(b);
            // 如果这两个点已经在同一个集合了，那么就不再合并了
            if (head1 == head2) {
                return;
            }
            // 将元素少的集合挂到元素多的集合下面
            Integer size1 = sizeMap.get(head1);
            Integer size2 = sizeMap.get(head2);
            Element<V> bigHead = size1 > size2 ? head1 : head2;
            Element<V> smallHead = bigHead == head1 ? head2 : head1;

            fatherMap.put(smallHead, bigHead);
            sizeMap.remove(smallHead);
            sizeMap.put(bigHead, size1 + size2);
        }

        /**
         * 判断两个元素是否属于同一个集合
         */
        public boolean isSameSet(V a, V b) {
            // 必须要先存在于元素集合中
            if (elementMap.containsKey(a) && elementMap.containsKey(b)) {
                Element<V> head1 = findHead(a);
                Element<V> head2 = findHead(b);
                return head1 == head2;
            }
            return false;
        }


        /**
         * 查找一个节点的头节点
         */
        private Element<V> findHead(V a) {
            Element<V> element = elementMap.get(a);
            // 不存在该节点
            if (element == null) {
                return null;
            }
            // 记录从该节点到他的头节点所走过的节点，压入栈中。用于后面优化链表长度
            Stack<Element<V>> path = new Stack<>();
            // 最顶层的头节点的父节点就是自己，所以只要找到父节点是自己本身的节点，那么该节点就是头节点
            while (element != fatherMap.get(element)) {
                path.push(element);
                element = fatherMap.get(element);
            }
            // 将沿途节点全部挂到新的头下（优化）
            while (!path.isEmpty()) {
                fatherMap.put(path.pop(), element);
            }
            return element;
        }
    }

}
